// Copyright (c) 2020 Presto Labs Pte. Ltd.
// Author: donggu

#pragma once

#include <memory>
#include <mutex>
#include <optional>
#include <queue>
#include <utility>

#include <boost/asio.hpp>
#include <boost/asio/high_resolution_timer.hpp>
#include <boost/system/error_code.hpp>
#include <sys/ioctl.h>

#include "coin2/base/log.h"
#include "coin2/base/ring.h"

/// LeakyBucket is enough for feed subscription.
/// For orders, please use TokenBucket.
class AsyncLeakyBucket : public std::enable_shared_from_this<AsyncLeakyBucket> {
 public:
  using Callback = std::function<void(boost::system::error_code)>;

  AsyncLeakyBucket(boost::asio::io_context::strand* strand, int64_t interval)
      : strand_(strand), interval_(interval), timer_(*strand) {
    CHECK_THROW(strand != nullptr);
    CHECK_GT_THROW(interval, 0);
  }
  AsyncLeakyBucket(boost::asio::io_context::strand* strand, int num, int64_t window_size)
      : strand_(strand), timer_(*strand) {
    CHECK_GT_THROW(num, 0);
    CHECK_GT_THROW(window_size, 0);

    // = floor(window_size / num)
    interval_ = (window_size - 1) / num + 1;
  }

  // TODO(jaewon): what if callback is deallocated while LeakyBucket is running?
  void PostAsync(Callback handler) {
    std::scoped_lock guard(mutex_);
    q_.push(handler);
    if (idle_) {
      idle_ = false;
      std::weak_ptr<AsyncLeakyBucket> self(shared_from_this());
      boost::asio::post(*strand_, [self] {
        if (auto s = self.lock()) {
          s->Leak();
        } else {
          LOG(ERROR) << "AsyncLeakyBucket is already destroyed";
        }
      });
    }
  }

  template <class CompletionToken>
  void PostCoro(CompletionToken&& token) {
    using Result = boost::asio::async_result<CompletionToken, void(boost::system::error_code)>;
    using Handler = typename Result::completion_handler_type;

    Handler handler(std::forward<CompletionToken>(token));
    Result result(handler);

    PostAsync(handler);

    result.get();
  }

 private:
  // this is blocking
  void Leak() {
    Callback callback = {};
    {
      std::scoped_lock guard(mutex_);
      if (q_.empty()) {
        // goes into idle.
        idle_ = true;
      } else {
        callback = q_.front();
        q_.pop();
        timer_.expires_after(std::chrono::nanoseconds(interval_));

        std::weak_ptr<AsyncLeakyBucket> self(shared_from_this());
        timer_.async_wait([self](const boost::system::error_code& error) {
          if (error) {
            LOG(ERROR) << "AsyncLeakyBucket error: " << error.message();
          }
          if (auto s = self.lock()) {
            s->Leak();
          } else {
            LOG(ERROR) << "AsyncLeakyBucket is already destroyed";
          }
        });
      }
    }
    // <- no lock at this point
    if (callback) {
      callback(boost::system::error_code());
    }
  }

  boost::asio::io_context::strand* strand_;
  bool idle_ = true;
  int64_t interval_;
  std::mutex mutex_;  // protects: idle_, q_, timer_
  std::queue<Callback> q_;
  boost::asio::high_resolution_timer timer_;
};

class TokenBucket {
  TokenBucket(boost::asio::io_context* strand, int num, int64_t window_size) {
    THROW() << "[TokenBucket] not implemented";
  }
};

// thread-safe
class SlidingRateLimiter {
 public:
  SlidingRateLimiter(int64_t window_size, size_t max_num)
      : window_size_(window_size), max_num_(max_num), ring_(max_num) {}

  size_t available(int64_t ts) {
    std::scoped_lock guard(mutex_);

    while (ring_.non_empty() && ring_.front() + window_size_ <= ts) {
      ring_.pop_front();
    }
    return max_num_ - ring_.size();
  }

  // returns 0 if no waiting
  int64_t expected_waiting(int64_t ts) {
    std::scoped_lock guard(mutex_);

    while (ring_.non_empty() && ring_.front() + window_size_ <= ts) {
      ring_.pop_front();
    }

    if (ring_.is_full()) {
      return ring_.front() + window_size_ - ts;
    }
    return 0;
  }

  // return true on success
  // return false on rate limit
  bool Take(int64_t ts) {
    std::scoped_lock guard(mutex_);

    while (ring_.non_empty() && ring_.front() + window_size_ <= ts) {
      ring_.pop_front();
    }

    if (ring_.is_full()) {
      return false;
    }
    ring_.push_back(ts);
    return true;
  }

  void DebugPrint() { ring_.DebugPrint(); }

 private:
  const int64_t window_size_;
  const size_t max_num_;
  coin2::base::Ring<int64_t> ring_;  // stores timestamps
  std::mutex mutex_;                 // protects: ring_
};
